# This is a sample Python script.

import cv2
# Press Shift+F10 to execute it or replace it with your code.
# Press Double Shift to search everywhere for classes, files, tool windows, actions, and settings.
import numpy as np
from PIL import Image


# 文档校正—透视变换
def perspective(src, src_pos):
    src_pos = np.array([src_pos[0], src_pos[1], src_pos[3], src_pos[2]])
    # 计算边长
    x_dist1 = np.linalg.norm(src_pos[1] - src_pos[0]).astype(int)
    x_dist2 = np.linalg.norm(src_pos[3] - src_pos[2]).astype(int)
    y_dist1 = np.linalg.norm(src_pos[2] - src_pos[0]).astype(int)
    y_dist2 = np.linalg.norm(src_pos[3] - src_pos[1]).astype(int)

    x_dist = np.int32((x_dist1 + x_dist2) / 2)
    y_dist = np.int32((y_dist1 + y_dist2) / 2)
    # 计算目标尺寸
    pic_scale = np.array([[0, 0], [x_dist, 0], [0, y_dist], [x_dist, y_dist]])

    # 转换坐标类型
    pos1 = np.float32(src_pos)
    pos2 = np.float32(pic_scale)

    # 获取变换矩阵
    # 原图中的四个角点pts1(对应好即可，左上、右上、左下、右下),与变换后矩阵位置pts2
    M = cv2.getPerspectiveTransform(pos1, pos2)
    # print(M)

    # 图像透视变换
    # cv2.warpPerspective(src, M， dsize=(cols, rows))
    # src: 原图
    # M： 一个3x3的变换矩阵
    # dsize: 输出图像的尺寸大小, 先指定(第一个参数是)col，再指定(第二个参数是)row
    result = cv2.warpPerspective(src, M, (x_dist, y_dist))

    return result


# 获取坐标点
# 原图中的四个角点pts1(对应好即可，左上、右上、右下、左下),与变换后矩阵位置pts2
def get_points(point_file_name):
    with open(point_file_name, "r") as label:
        label_str = label.readline()
        points = label_str.split(" ")
    return np.array(points[:8], dtype=np.int32).reshape(4, 2)


def loadLabel(path, index):
    mask_points = []
    if path is not "":
        point_path = "%s/train_%s.txt" % (path, index)
        mask_points = get_points(point_path)
    return mask_points


if __name__ == '__main__':
    for i in range(1, 1001):
        src = cv2.imread('data/train_imgs/train_%s.jpg' % i)
        pos = loadLabel('data/labels', i)
        target_pic = perspective(src, pos)

        x, y = 120, 30
        fig = cv2.resize(target_pic, (x, y))
        img = Image.fromarray(fig)
        image_path = 'data/result/%04d.jpg' % i
        img.save(image_path)
        print('Image_%04d' % i, '-->', image_path)

        step = x / 5
        for j in range(1, 6):
            box = ((j - 1) * step, 0., j * step, y)
            # print(box)
            sub_img = img.crop(box)
            sub_img_path = 'data/result5/%04d_%d.jpg' % (i, j)
            sub_img.save(sub_img_path)
            print('SubImage_%d' % j, '-->', sub_img_path)
